/****************************************************************************
*
* Description: PE file format reader
* Module: 'pe.h'
* Created at: 24.06.06 17:11:38 UTC
*
* $Revision$
*
****************************************************************************/
#pragma once

#if !defined(__STORAGE_PE_PE32_H__)
#define __STORAGE_PE_PE32_H__

#include <string>
#include <common/defs.h>
#include <storage/basic_file.h>

namespace storage {
namespace pe {

#include <storage/pe/pedefs.h>

/* I think that it can be called "Pe32Loader".
 * Use it to map PE32 to memory.  */
struct pe32 :
    public basic_file
{
	/* Default ctor.  */
    pe32():
        basic_file()
    {}

	/* Ctor. Name specified.  */
    pe32(std::string & name):
        basic_file(name)
    {}

	/* Ctor. Data provided.  */
    pe32(bytes_t & data):
        basic_file(data)
    {}

	/* Ctor. Full defined version.  */
    pe32(std::string & name, bytes_t & data):
        basic_file(name, data)
    {}

	/* Dtor.  */
	virtual ~pe32()
	{}

	/* --------- source_t interface ----------  */
    /* virtual void read() ...  */

	/* Check that data is valid for this source.  */
    virtual bool is_valid()
    {
        PIMAGE_DOS_HEADER       pdh;
        PIMAGE_NT_HEADERS       pnh;

        if( data().size() < sizeof(IMAGE_DOS_HEADER) )
            return false;

        pdh = reinterpret_cast<PIMAGE_DOS_HEADER>(&data()[0]);
        if( pdh->e_magic != IMAGE_DOS_SIGNATURE )
            return false;

        if( data().size() < pdh->e_lfanew + sizeof(IMAGE_NT_HEADERS) )
            return false;

        pnh = reinterpret_cast<PIMAGE_NT_HEADERS>(&data()[pdh->e_lfanew]);
        if( pnh->Signature != IMAGE_NT_SIGNATURE )
            return false;

        if( pnh->FileHeader.SizeOfOptionalHeader < sizeof(IMAGE_OPTIONAL_HEADER) )
            return false;

        return true;
    }

	/* Routine maps provided file to memory.  */
	virtual void map_to(common::memory::memory_manager_t & mem_mgr)
	{
		headers		h(*const_cast<pe32 *>(this));
		sections_t	s(*const_cast<pe32 *>(this));
		sections_t::container_t::iterator	b = s.sections().begin(),
											e = s.sections().end(),
											i;
		/* for each section in the file.  */
		for (i = b; i != e; ++i)
		{
			/* Create separate memory block.  */
			common::memory::memory_manager_t::region_t & r =
				mem_mgr.create_block(h.poh->ImageBase + i->vaddr, i->vsize);
			/* Copy section data to memory block.  */
			// r.label().assign(i->name);
			memcpy(&r[0], &data()[i->raddr], std::min(i->vsize, i->rsize));
		}
	}

	/* Virtual Address to Raw Address.  */
	uint32 va2ra(uint32 ra)
	{
		sections_t s(*const_cast<pe32 *>(this));
		sections_t::container_t::iterator	le = s.sections().end(),
											lb = s.sections().begin(),
											li = le - 1;
		
		if (ra > li->vaddr + li->vsize) return -1;
		for (li = lb; li != le && ra >= li->vaddr; ++li);
		--li;
		return li->raddr + (ra - li->vaddr);
	}

	/* Return list of file entry points.  */
	virtual void get_entry_points(addrs_t & eps)
	{
		headers h(*const_cast<pe32 *>(this));

		/* Module entry point should be first.  */
		if (h.poh->AddressOfEntryPoint)
			eps.push_back(h.poh->ImageBase + h.poh->AddressOfEntryPoint);

		/* TODO: Possible this code should be transfered to external set of pe tools,
		* so the additional information about export can be collected.  */
		/* Process export table.  */
		if (h.poh->NumberOfRvaAndSizes > IMAGE_DIRECTORY_ENTRY_EXPORT
			&& h.poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
			&& h.poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].Size)
		{
			/* We need raw address in the file instead of RVA.  */
			uint32 ra = va2ra(h.poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress);

			/* FIXME: Check that ra != -1.  */
			PIMAGE_EXPORT_DIRECTORY pIED =
				reinterpret_cast<PIMAGE_EXPORT_DIRECTORY>(&data()[ra]);
			uint32 fra = va2ra(pIED->AddressOfFunctions); /* Functions RA.  */
			puint32 fns = reinterpret_cast<puint32>(&data()[fra]);

			for(uint32 i = 0; i < pIED->NumberOfFunctions; ++i)
			{
				/* Array should be continuous, so ms specifies
				 * that missed fns has 0 as RVA.  */
				if (fns[i])
					eps.push_back(h.poh->ImageBase + fns[i]);
			}
		}
	}
	/* ------- source_t interface end --------  */

	/* ---------- sink_t interface -----------  */
    /* virtual void write() ...  */
	/* -------- sink_t interface end ---------  */

    /* PE headers container.  */
    struct headers {
		/* Ctor. File should be specified.  */
        explicit headers(pe32 & image)
        {
            update(image);
        }

		/* Update invalidated headers.  */
        void update(pe32 & image)
        {
            pdh = NULL; pnh = NULL; pfh = NULL; poh = NULL; psh = NULL;
            if( image.is_valid() )
            {
                pdh = reinterpret_cast<PIMAGE_DOS_HEADER>(&image.data()[0]);
                pnh = reinterpret_cast<PIMAGE_NT_HEADERS>(&image.data()[pdh->e_lfanew]);
                pfh = &pnh->FileHeader;
                poh = &pnh->OptionalHeader;
                psh = reinterpret_cast<PIMAGE_SECTION_HEADER>(
					reinterpret_cast<puint8>(poh) + pnh->FileHeader.SizeOfOptionalHeader);
            }
        }

		/* Data members.  */
        PIMAGE_DOS_HEADER       pdh;
        PIMAGE_NT_HEADERS       pnh;
        PIMAGE_FILE_HEADER      pfh;
        PIMAGE_OPTIONAL_HEADER  poh;
        PIMAGE_SECTION_HEADER   psh;

    }; /* End of headers.  */

	/* PE section object.  */
	struct section_t {

		/* Default ctor.  */
		section_t()
		{
			vaddr = vsize = raddr = rsize = attr = 0;
		}

		/* Copy ctor required for STL containers.  */
		section_t(const section_t & o)
		{
			name.assign(o.name.c_str());
			vaddr = o.vaddr;
			vsize = o.vsize;
			raddr = o.raddr;
			rsize = o.rsize;
			attr  = o.attr;
		}

		/* Data members.  */
		std::string name;   /* Section name - 8 chars max.  */
		uint32       vaddr;
		uint32       vsize;  /* Virtual size - in memory.  */
		uint32       raddr;  /* Raw addr - offset in file.  */
		uint32       rsize;  /* Raw size - size in file.  */
		uint32       attr;   /* Section attributes.  */

	}; /* End of section_t.  */

	/* PE sections set.  */
	struct sections_t {

		typedef std::vector<section_t>    container_t;

		/* Ctor. File image should be specified.  */
		explicit sections_t(pe32 & image)
		{
			update(image);
		}

		/* Update invalidated section list for provided file image.  */
		void update(pe32 & image)
		{
			headers   h(image);
			sections_.clear();
			for(int i = 0; i < h.pfh->NumberOfSections; ++i)
			{
				section_t	s;
				s.name.assign(&h.psh[i].Name[0], &h.psh[i].Name[IMAGE_SIZEOF_SHORT_NAME]);
				s.vaddr     = h.psh[i].VirtualAddress;
				s.vsize     = h.psh[i].Misc.VirtualSize;
				s.raddr     = h.psh[i].PointerToRawData;
				s.rsize     = h.psh[i].SizeOfRawData;
				s.attr      = h.psh[i].Characteristics;

				sections_.push_back(s);
			}
		}

		/* Direct access to sections list.  */
		container_t & sections()
		{
			return sections_;
		}

		/* Data members.  */
		container_t	sections_;

	}; /* End of sections_t.  */

}; /* End of file.  */

} /* End of pe.  */
} /* End of storage.  */

#endif  /* End of __STORAGE_PE_PE32_H__.  */
